import os,yaml
import sys
import json
from pathlib import Path
from PyQt5 import QtGui,QtCore,QtWidgets
from PyQt5.QtGui import QIcon, QPixmap
from PyQt5.QtCore import Qt
from ultralytics import YOLO
from threading import Thread
from PyQt5.QtCore import *
from PyQt5.QtWidgets import *
from PyQt5.QtWidgets import QFileDialog,QMessageBox
import photoshower,Tensorboradshow
from test import QMainWindow_new

class Ui_MainWindow(object):
    font1 = QtGui.QFont("Adobe 黑体 Std R")
    font1.setPixelSize(32)
    font2 = QtGui.QFont("Adobe 黑体 Std R")
    font2.setPixelSize(24)
    font3 = QtGui.QFont("Adobe 黑体 Std R")
    font3.setPixelSize(16)

    singal = pyqtSignal(int)
    singal_2 = pyqtSignal(int)
    singal_3 = pyqtSignal(int)
    singal_4= pyqtSignal(dict)

    # singals = [None,singal,singal_2,singal_3,singal_4]
    device_dict = {"cpu":0, 0:1, "tpu":2}

    def __init__(self):
        with open('runin_train.json',encoding='utf-8') as file:
            self.D=json.load(file)

    def setupUi(self, MainWindow):
        MainWindow.setObjectName("MainWindow")
        MainWindow.resize(1200, 650)

        label_backdrop = QLabel(MainWindow)
        label_backdrop.resize(1200, 650)
        label_backdrop.setStyleSheet("background-color: rgb(240, 250, 250);")

        label_title = QLabel(MainWindow)
        label_title.resize(1200, 50)
        label_title.setText("    训练模型")
        label_title.setFont(QtGui.QFont("Adobe 黑体 Std R", 16))
        label_title.setStyleSheet("background-color: rgb(50, 200, 200);")

        label_T = QLabel(MainWindow)
        label_T.resize(40, 40)
        label_T.move(5, 5)
        label_T.setPixmap(QPixmap(r"ui\label.jpg").scaled(label_T.size(), Qt.KeepAspectRatio, Qt.SmoothTransformation))

        button_close = QPushButton(MainWindow)
        button_close.setStyleSheet("QPushButton {border: none;}")
        button_close.setIcon(QIcon(r"ui\close.png"))
        button_close.resize(50, 50)
        button_close.move(1150, 0)
        button_close.clicked.connect(MainWindow.close)

        button_min = QPushButton(MainWindow)
        button_min.setStyleSheet("QPushButton {border: none;}")
        button_min.setIcon(QIcon(r"ui\min.png"))
        button_min.resize(50, 50)
        button_min.move(1100, 0)
        button_min.clicked.connect(MainWindow.showMinimized)

        # self.centralwidget = QtWidgets.QWidget(MainWindow)
        # self.centralwidget.setObjectName("centralwidget")
        # self.horizontalLayout = QtWidgets.QHBoxLayout(self.centralwidget)
        # self.horizontalLayout.setObjectName("horizontalLayout")
# —————————————————————————————————————————————————————————————————————

        self.label_set = QtWidgets.QLabel(MainWindow)
        self.label_set.setGeometry(QtCore.QRect(20, 60, 200, 30))
        self.label_set.setObjectName("label_15")
# ————————————————————————————————————————————————————————————————————————————
        self.label_model = QtWidgets.QLabel(MainWindow)
        self.label_model.setGeometry(QtCore.QRect(30, 120, 120, 30))
        self.label_model.setObjectName("label_7")

        self.lineEdit_model = QtWidgets.QLineEdit(MainWindow)
        self.lineEdit_model.setText(self.D['model'])
        self.lineEdit_model.setGeometry(QtCore.QRect(150, 120, 300, 30))
        self.lineEdit_model.setObjectName("lineEdit_4")

        self.toolButton_model = QtWidgets.QToolButton(MainWindow)
        self.toolButton_model.setGeometry(QtCore.QRect(470, 120, 60, 30))
        self.toolButton_model.setObjectName("toolButton")
# ————————————————————————————————————————————————————————————————————————————
        self.label_yaml = QtWidgets.QLabel(MainWindow)
        self.label_yaml.setGeometry(QtCore.QRect(30, 170, 170, 30))
        self.label_yaml.setObjectName("label_8")

        self.lineEdit_yaml = QtWidgets.QLineEdit(MainWindow)
        self.lineEdit_yaml.setText(self.D['yaml'])
        self.lineEdit_yaml.setGeometry(QtCore.QRect(150, 170, 300, 30))
        self.lineEdit_yaml.setObjectName("lineEdit_5")

        self.toolButton_yaml = QtWidgets.QToolButton(MainWindow)
        self.toolButton_yaml.setGeometry(QtCore.QRect(470, 170, 60, 30))
        self.toolButton_yaml.setObjectName("toolButton_2")
# ————————————————————————————————————————————————————————————————————————————
        self.label_path = QtWidgets.QLabel(MainWindow)
        self.label_path.setGeometry(QtCore.QRect(30, 220, 170, 30))
        self.label_path.setObjectName("label_11")

        self.lineEdit_path = QtWidgets.QLineEdit(MainWindow)
        self.lineEdit_path.setText(self.D['Path'])
        self.lineEdit_path.setGeometry(QtCore.QRect(150, 220, 300, 30))
        self.lineEdit_path.setObjectName("lineEdit_8")

        self.toolButton_path = QtWidgets.QToolButton(MainWindow)
        self.toolButton_path.setGeometry(QtCore.QRect(470, 220, 60, 30))
        self.toolButton_path.setObjectName("toolButton_5")
# ————————————————————————————————————————————————————————————————————————————
        self.label_epochs = QtWidgets.QLabel(MainWindow)
        self.label_epochs.setGeometry(QtCore.QRect(30, 300, 170, 30))
        self.label_epochs.setObjectName("label_5")

        self.lineEdit_epochs = QtWidgets.QLineEdit(MainWindow)
        self.lineEdit_epochs.setText(str(self.D['epochs']))
        self.lineEdit_epochs.setGeometry(QtCore.QRect(170, 300, 80, 30))
        self.lineEdit_epochs.setObjectName("lineEdit_2")
# ————————————————————————————————————————————————————————————————————————————
        self.label_wait = QtWidgets.QLabel(MainWindow)
        self.label_wait.setGeometry(QtCore.QRect(310, 300, 120, 30))
        self.label_wait.setObjectName("label_10")

        self.lineEdit_wait = QtWidgets.QLineEdit(MainWindow)
        self.lineEdit_wait.setText(str(self.D['wait']))
        self.lineEdit_wait.setGeometry(QtCore.QRect(450, 300, 80, 30))
        self.lineEdit_wait.setObjectName("lineEdit_6")
# ————————————————————————————————————————————————————————————————————————————
        self.label_imgsz = QtWidgets.QLabel(MainWindow)
        self.label_imgsz.setGeometry(QtCore.QRect(30, 350, 120, 30))
        self.label_imgsz.setObjectName("label_14")

        self.lineEdit_imgsz = QtWidgets.QLineEdit(MainWindow)
        self.lineEdit_imgsz.setText(str(self.D['imgsz']))
        self.lineEdit_imgsz.setGeometry(QtCore.QRect(170, 350, 80, 30))
        self.lineEdit_imgsz.setObjectName("lineEdit_9")
# ————————————————————————————————————————————————————————————————————————————
        self.label_batch = QtWidgets.QLabel(MainWindow)
        self.label_batch.setGeometry(QtCore.QRect(310, 350, 120, 30))
        self.label_batch.setObjectName("label_13")

        self.lineEdit_batch = QtWidgets.QLineEdit(MainWindow)
        self.lineEdit_batch.setText(str(self.D['batch']))
        self.lineEdit_batch.setGeometry(QtCore.QRect(450, 350, 80, 30))
        self.lineEdit_batch.setObjectName("lineEdit_7")
# ————————————————————————————————————————————————————————————————————————————
        self.label_name = QtWidgets.QLabel(MainWindow)
        self.label_name.setGeometry(QtCore.QRect(30, 400, 120, 30))
        self.label_name.setObjectName("label_6")

        self.lineEdit_name = QtWidgets.QLineEdit(MainWindow)
        self.lineEdit_name.setText(self.D['name'])
        self.lineEdit_name.setGeometry(QtCore.QRect(170, 400, 80, 30))
        self.lineEdit_name.setObjectName("lineEdit_3")
# ————————————————————————————————————————————————————————————————————————————
        self.label_device = QtWidgets.QLabel(MainWindow)
        self.label_device.setGeometry(QtCore.QRect(310, 400, 120, 30))
        self.label_device.setObjectName("label_12")

        self.comboBox_device = QtWidgets.QComboBox(MainWindow)
        self.comboBox_device.setGeometry(QtCore.QRect(450, 400, 80, 30))
        font = QtGui.QFont()
        font.setPointSize(10)
        self.comboBox_device.setFont(font)
        self.comboBox_device.setObjectName("comboBox")
        self.comboBox_device.addItems(["cpu", "gpu", "tpu"])
        self.comboBox_device.setCurrentIndex(self.device_dict[self.D["device"]])
# ————————————————————————————————————————————————————————————————————————————
        self.label_box = QtWidgets.QLabel(MainWindow)
        self.label_box.setGeometry(QtCore.QRect(30, 450, 120, 30))

        self.lineEdit_box = QtWidgets.QLineEdit(MainWindow)
        self.lineEdit_box.setText(str(self.D['box']))
        self.lineEdit_box.setGeometry(QtCore.QRect(170, 450, 80, 30))
# ————————————————————————————————————————————————————————————————————————————
        self.label_cls = QtWidgets.QLabel(MainWindow)
        self.label_cls.setGeometry(QtCore.QRect(310, 450, 120, 30))

        self.lineEdit_cls = QtWidgets.QLineEdit(MainWindow)
        self.lineEdit_cls.setText(str(self.D['cls']))
        self.lineEdit_cls.setGeometry(QtCore.QRect(450, 450, 80, 30))
# ————————————————————————————————————————————————————————————————————————————
        self.label_new = QtWidgets.QLabel(MainWindow)
        self.label_new.setGeometry(QtCore.QRect(30, 520, 200, 30))
        self.label_new.setObjectName("label_9")

        self.checkBox_new = QtWidgets.QCheckBox(MainWindow)
        self.checkBox_new.setGeometry(QtCore.QRect(220, 520, 80, 30))
        self.checkBox_new.setObjectName("checkBox")
        self.checkBox_new.setChecked(self.D['exist_ok'])
# ————————————————————————————————————————————————————————————————————————————————
        self.pushButton_look_img = QtWidgets.QPushButton(MainWindow)
        self.pushButton_look_img.setGeometry(QtCore.QRect(350, 520, 180, 35))
        self.pushButton_look_img.setFont(self.font2)
        self.pushButton_look_img.setObjectName("pushButton_3")
# ————————————————————————————————————————————————————————————————————————————————
        self.pushButton_set = QtWidgets.QPushButton(MainWindow)
        self.pushButton_set.setGeometry(QtCore.QRect(30, 570, 120, 35))
        self.pushButton_set.setFont(self.font2)
        self.pushButton_set.setObjectName("pushButton_2")
# ————————————————————————————————————————————————————————————————————————————————
        self.pushButton_train = QtWidgets.QPushButton(MainWindow)
        self.pushButton_train.setGeometry(QtCore.QRect(180, 570, 120, 35))
        self.pushButton_train.setFont(self.font2)
        self.pushButton_train.setObjectName("pushButton")
# ————————————————————————————————————————————————————————————————————————————————
        self.pushButton_look_train = QtWidgets.QPushButton(MainWindow)
        self.pushButton_look_train.setGeometry(QtCore.QRect(350, 570, 180, 35))
        self.pushButton_look_train.setFont(self.font2)
        self.pushButton_look_train.setObjectName("pushButton_4")
        # self.horizontalLayout.addWidget(self.frame)
# ————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————————

        self.label_book = QtWidgets.QLabel(MainWindow)
        self.label_book.setGeometry(QtCore.QRect(600, 60, 200, 30))
        self.label_book.setMaximumSize(QtCore.QSize(150, 30))
        self.label_book.setObjectName("label")

        self.textBrowser_book = QtWidgets.QTextBrowser(MainWindow)
        self.textBrowser_book.setGeometry(QtCore.QRect(620,120, 530, 370))
        self.textBrowser_book.setObjectName("textBrowser")
# —————————————————————————————————————————————————————————————————————
        self.label_now_train = QtWidgets.QLabel(MainWindow)
        self.label_now_train.setGeometry(QtCore.QRect(620, 525, 200, 30))
        self.label_now_train.setObjectName("label_2")

        self.progressBar_now_train = QtWidgets.QProgressBar(MainWindow)
        self.progressBar_now_train.setGeometry(QtCore.QRect(820, 525, 350, 30))
        self.progressBar_now_train.setCursor(QtGui.QCursor(QtCore.Qt.ArrowCursor))
        self.progressBar_now_train.setProperty("value", 0)
        self.progressBar_now_train.setObjectName("progressBar")
# —————————————————————————————————————————————————————————————————————
        self.label_all_train = QtWidgets.QLabel(MainWindow)
        self.label_all_train.setGeometry(QtCore.QRect(620, 575, 200, 30))
        self.label_all_train.setObjectName("label_3")

        self.progressBar_all_train = QtWidgets.QProgressBar(MainWindow)
        self.progressBar_all_train.setGeometry(QtCore.QRect(820, 575, 350, 30))
        self.progressBar_all_train.setProperty("value", 0)
        self.progressBar_all_train.setObjectName("progressBar_2")
# —————————————————————————————————————————————————————————————————————
#         self.horizontalLayout.addWidget(self.frame_2)

        # MainWindow.setCentralWidget(self.centralwidget)
        self.menubar = QtWidgets.QMenuBar(MainWindow)
        self.menubar.setGeometry(QtCore.QRect(0, 0, 947, 30))
        self.menubar.setObjectName("menubar")
        MainWindow.setMenuBar(self.menubar)
        self.retranslateUi(MainWindow)
        QtCore.QMetaObject.connectSlotsByName(MainWindow)

    def retranslateUi(self, MainWindow):
        _translate = QtCore.QCoreApplication.translate
        MainWindow.setWindowTitle(_translate("MainWindow", "训练模型"))
        self.label_set.setText(_translate("MainWindow", "<html><head/><body><p>基础设置:</span></p></body></html>"))
        self.toolButton_model.setText(_translate("MainWindow", "file"))
        self.toolButton_yaml.setText(_translate("MainWindow", "file"))
        self.toolButton_path.setText(_translate("MainWindow", "dir"))
        self.label_epochs.setText(_translate("MainWindow", "<html><head/><body><p>训练轮数:</span></p></body></html>"))
        self.label_name.setText(_translate("MainWindow", "<html><head/><body><p>模型名称:</span></p></body></html>"))
        self.label_model.setText(_translate("MainWindow", "<html><head/><body><p>模型文件:</span></p></body></html>"))
        self.label_yaml.setText(_translate("MainWindow", "<html><head/><body><p>数据文件:</span></p></body></html>"))
        self.label_path.setText(_translate("MainWindow", "<html><head/><body><p>保存路径:</span></p></body></html>"))
        self.label_device.setText(_translate("MainWindow", "<html><head/><body><p>训练设备:</span></p></body></html>"))
        self.label_box.setText(_translate("MainWindow","<html><head/><body><p>边框损失:</span></p></body></html>"))
        self.label_cls.setText(_translate("MainWindow","<html><head/><body><p>类别损失:</span></p></body></html>"))
        self.checkBox_new.setText(_translate("MainWindow", "开启"))
        self.label_new.setText(_translate("MainWindow", "<html><head/><body><p>覆盖同名实验:</span></p></body></html>"))
        self.pushButton_set.setText(_translate("MainWindow", "保存设置"))
        self.pushButton_train.setText(_translate("MainWindow", "开始训练"))
        self.label_wait.setText(_translate("MainWindow", "<html><head/><body><p>早停轮数:</span></p></body></html>"))
        self.label_batch.setText(_translate("MainWindow", "<html><head/><body><p>批量大小:</span></p></body></html>"))
        self.label_imgsz.setText(_translate("MainWindow", "<html><head/><body><p>图片大小:</span></p></body></html>"))
        self.pushButton_look_img.setText(_translate("MainWindow", "查看数据图片"))
        self.pushButton_look_train.setText(_translate("MainWindow", "查看训练结果"))
        self.label_book.setText(_translate("MainWindow", "<html><head/><body><p>日志记录</span></p></body></html>"))
        self.label_now_train.setText(_translate("MainWindow", "<html><head/><body><p>现批次训练进度:</span></p></body></html>"))
        self.label_all_train.setText(_translate("MainWindow", "<html><head/><body><p>总批次训练进度:</span></p></body></html>"))

        self.label_set.setFont(self.font1)
        self.label_book.setFont(self.font1)
        self.label_epochs.setFont(self.font2)
        self.label_name.setFont(self.font2)
        self.label_model.setFont(self.font2)
        self.label_yaml.setFont(self.font2)
        self.label_path.setFont(self.font2)
        self.label_device.setFont(self.font2)
        self.label_box.setFont(self.font2)
        self.label_cls.setFont(self.font2)
        self.label_new.setFont(self.font2)
        self.label_wait.setFont(self.font2)
        self.label_batch.setFont(self.font2)
        self.label_imgsz.setFont(self.font2)
        self.label_now_train.setFont(self.font2)
        self.label_all_train.setFont(self.font2)

class MyWindow(QMainWindow_new, Ui_MainWindow):
    def __init__(self):
        super(MyWindow, self).__init__()
        self.setupUi(self)

        self.toolButton_model.clicked.connect(self.select_model)
        self.toolButton_yaml.clicked.connect(self.select_yaml)
        self.toolButton_path.clicked.connect(self.select_savePath)
        self.pushButton_set.clicked.connect(self.save_all)
        self.pushButton_train.clicked.connect(self.start_train)
        self.pushButton_look_img.clicked.connect(self.ImageViewer)
        self.pushButton_look_train.clicked.connect(self.ProcessView)

    def select_model(self):
        fname,_=QFileDialog.getOpenFileName(self,'Open file','./','Image files(*.pt)')
        self.lineEdit_model.setText(fname)

    def select_yaml(self):
        fname , _=QFileDialog.getOpenFileName(self,'Open file','./','Image files(*.yaml)')
        self.lineEdit_yaml.setText(fname)

    def select_savePath(self):
        directory = QtWidgets.QFileDialog.getExistingDirectory(None, "选择文件夹", "D:/")
        self.lineEdit_path.setText(directory)


    def save_all(self):
        self.D['model']=self.lineEdit_model.text()
        self.D['yaml']=self.lineEdit_yaml.text()
        self.D['Path']=self.lineEdit_path.text()
        self.D['epochs']=int(self.lineEdit_epochs.text())
        self.D['name']=self.lineEdit_name.text()
        self.D['wait']=int(self.lineEdit_wait.text())
        self.D['batch']=int(self.lineEdit_batch.text())
        self.D['imgsz']=int(self.lineEdit_imgsz.text())
        self.D['box']=float(self.lineEdit_box.text())
        self.D['cls'] = float(self.lineEdit_cls.text())
        self.D['exist_ok']=self.checkBox_new.isChecked()
        if self.comboBox_device.currentText()== 'cpu':
            self.D['device']=self.comboBox_device.currentText()
        elif self.comboBox_device.currentText()== 'gpu':
            self.D['device'] = 0
        elif self.comboBox_device.currentText()== 'tpu':
            self.D['device']= 'tpu'
        else:
            self.D['device']=-1
        sub_dir = 'Ultralytics'
        path = Path.home() / 'AppData' / 'Roaming' / sub_dir / 'settings.yaml'
        try:
            with open(path,'r') as file:
                confign = yaml.safe_load(file)
                confign['runs_dir'] = self.D['Path']
                confign['weights_dir'] = self.D['Path']
            with open(path,'w') as file:
                yaml.safe_dump(confign, file)
        except:
            pass
        with open('runin_train.json', 'w') as file:
            json.dump(self.D, file, indent=4)
        QMessageBox.information(self,"操作成功","保存成功")

    # 当前批次训练进度传递
    def temp_progressBar(self, value):
        self.progressBar_now_train.setValue(value)
    # 总批次训练进度传递
    def temp_progressBar_2(self, value):
        self.progressBar_all_train.setValue(value)
        if value==self.D["epochs"]:
            QMessageBox.information(self, "训练完成", "保存到" + self.D["Path"] + "/runs/detect/" + self.D["name"])

    def getpBar(self,value):
        self.V = value
        self.progressBar_now_train.setRange(0, value)

    def getmetrics(self, value):
        self.textBrowser_book.append('第' + str(self.i) + '轮\n')
        self.textBrowser_book.append(('metrics/precision(B):' +'\t'+ str(value['metrics/precision(B)'])+'\n'))
        self.textBrowser_book.append(('metrics/recall(B):' +'\t\t'+ str(value['metrics/recall(B)'])+'\n'))
        self.textBrowser_book.append(('metrics/mAP50(B)' +'\t\t'+ str(value['metrics/mAP50(B)'])+'\n'))
        self.textBrowser_book.append(('metrics/mAP50-95(B)' +'\t\t'+ str(value['metrics/mAP50-95(B)'])+'\n'))
        self.textBrowser_book.append(('val/box_loss' +'\t\t'+ str(value['val/box_loss'])+'\n'))
        self.textBrowser_book.append(('val/cls_loss' +'\t\t'+ str(value['val/cls_loss'])+'\n'))
        self.textBrowser_book.append(('val/dfl_loss' +'\t\t'+ str(value['val/dfl_loss'])+'\n'))
        self.i += 1


    def ImageViewer(self):
        self.photo_page = photoshower.Ui_photoshower()
        try:
            with open(self.D["yaml"]) as file:
                config = yaml.safe_load(file)
                lpath = config["path"]
                if config["train"] != None:
                    self.photo_page.datasets_train = lpath+"/"+config["train"]
                if config["val"] != None:
                    self.photo_page.datasets_val = lpath+"/"+config["val"]
                if config["test"] != None:
                    self.photo_page.datasets_test = lpath + "/" + config["test"]
            self.photo_page.show()
        except:
            QMessageBox.warning(self,"错误", "文件不存在")

    def ProcessView(self):
        self.ten=Tensorboradshow.Ui_tensorboradshow()
        self.ten.show()

    def start_train(self):
        self.i = 1
        self.textBrowser_book.clear()
        self.progressBar_all_train.setRange(0, self.D['epochs'])
        self.singal.emit(0)
        try:
            target = self.begin_train
            Thread(target=target).start()
        except:
            QMessageBox.warning(self,"运行失败","请检查输入")

    def begin_train(self):
        self.singal.connect(self.temp_progressBar_2)
        self.singal_2.connect(self.temp_progressBar)
        self.singal_3.connect(self.getpBar)
        self.singal_4.connect(self.getmetrics)
        model = YOLO(self.D['model'])
        model.train(data=self.D['yaml'],
                    save_dir=self.D['Path'],
                    epochs=self.D['epochs'],
                    imgsz=self.D['imgsz'],
                    patience=self.D['wait'],
                    batch=self.D['batch'],
                    name=self.D['name'],
                    device=self.D['device'],
                    box=self.D["box"],
                    cls=self.D["cls"],
                    exist_ok=self.D['exist_ok'],
                    workers=0,
                    singal=self.singal,
                    singal_2=self.singal_2,
                    singal_3=self.singal_3,
                    singal_4=self.singal_4)
        self.singal_2.emit(self.V)
        self.singal.emit(self.D['epochs'])
        os.system("tensorboard --logdir="+self.D["Path"]+"/detect/"+self.D["name"])

if __name__ == "__main__":
    QCoreApplication.setAttribute(QtCore.Qt.AA_EnableHighDpiScaling)
    app = QApplication([])
    stats = MyWindow()
    stats.show()
    app.exec_()
    sys.exit()